import uuid
from typing import List, Dict
from datetime import datetime
from core.utils.inner_request_util import inner_send_chat
import threading


class Message:
    def __init__(self, role: str, content: str = None, uniq_id: str = None, tool_calls = None, tool_call_id=None):
        self.uniq_id = uniq_id if uniq_id is not None else str(uuid.uuid4())
        self.role = role
        self.content = content
        self.tool_calls = tool_calls
        self.tool_call_id = tool_call_id


class Dialogue:
    def __init__(self, conn):
        self.dialogue: List[Message] = []
        # 获取当前时间
        self.current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

        self.conn = conn
        self.chat_history_id = None
        self.chat_sequence = 0
        self.chat_sequence_lock = threading.Lock()

    def put(self, message: Message):
        self.dialogue.append(message)
        try:
            self.send_chat(
                message.role,
                message.tool_calls[0]['function']['name'] if message.content is None else message.content,
                None,
                None,
                None,
                None,
                None,
                None,
                self.get_and_add_chat_sequence(),
            )
        except Exception as e:
            self.conn.logger.bind(tag=self.conn.TAG).error(f"发生错误: {e}")

    def getMessages(self, m, dialogue):
        if m.tool_calls is not None:
            dialogue.append({"role": m.role, "tool_calls": m.tool_calls})
        elif m.role == "tool":
            dialogue.append({"role": m.role, "tool_call_id": m.tool_call_id, "content": m.content})
        else:
            dialogue.append({"role": m.role, "content": m.content})

    def get_llm_dialogue(self) -> List[Dict[str, str]]:
        dialogue = []
        for m in self.dialogue:
            self.getMessages(m, dialogue)
        return dialogue

    def update_system_message(self, new_content: str):
        """更新或添加系统消息"""
        # 查找第一个系统消息
        system_msg = next((msg for msg in self.dialogue if msg.role == "system"), None)
        if system_msg:
            system_msg.content = new_content
        else:
            self.put(Message(role="system", content=new_content))

    def get_llm_dialogue_with_memory(self, memory_str: str = None) -> List[Dict[str, str]]:
        if memory_str is None or len(memory_str) == 0:
            return self.get_llm_dialogue()
        
        # 构建带记忆的对话
        dialogue = []
        
        # 添加系统提示和记忆
        system_message = next(
            (msg for msg in self.dialogue if msg.role == "system"), None
        )


        if system_message:
            enhanced_system_prompt = (
                f"{system_message.content}\n\n"
                f"相关记忆：\n{memory_str}"
            )
            dialogue.append({"role": "system", "content": enhanced_system_prompt})

        # 添加用户和助手的对话
        for m in self.dialogue:
            if m.role != "system":  # 跳过原始的系统消息
                self.getMessages(m, dialogue)

        return dialogue

    def get_and_add_chat_sequence(self):
        with self.chat_sequence_lock:
            self.chat_sequence += 1
        return self.chat_sequence

    def send_chat(self, role, content, prompt_tokens, total_tokens, completion_tokens, prompt_ms, total_ms, completion_ms, sequence):
        if self.chat_history_id is None:
            future = self.conn.executor.submit(
                inner_send_chat,
                self.chat_history_id,
                self.conn.headers.get("device-id"),
                self.conn.headers.get("face-id"),
                role,
                content,
                prompt_tokens,
                total_tokens,
                completion_tokens,
                prompt_ms,
                total_ms,
                completion_ms,
                sequence,
                self.conn.logger
            )
            self.chat_history_id = future.result()
        else:
            self.conn.executor.submit(
                inner_send_chat,
                self.chat_history_id,
                self.conn.headers.get("device-id"),
                self.conn.headers.get("face-id"),
                role,
                content,
                prompt_tokens,
                total_tokens,
                completion_tokens,
                prompt_ms,
                total_ms,
                completion_ms,
                sequence,
                self.conn.logger
            )
